#include <TStopwatch.h>
#include <CpLog.h>

#include "CpFitter.h"
#include "CpPolynom1D.h"
#include "CpPolynom2D.h"
#include "CpPolynom3D.h"
#include <CpNames.h>

ClassImp(CpFitter);

//______________________________________________________________________________
CpFitter::CpFitter(): TObject(),
    fOutFile(0),
    fOutputDirName("/tmp"),
    fOutputFileName("defaultOutFile.root"),
    fPolynomGradeNumber(1),
    fMagneticAxisIndex(0) {
  CpDebug(CpLog::kDebug + 5, "<-");
  Int_t i;
  for (i = 0;i < 3;i++) {
    fMinMaxBin[i][0] = 0.0;
    fMinMaxBin[i][1] = 1.0;
    fMinMaxBin[i][2] = 0.0;
  }

  SetActivePositionAxis();
  CpDebug(CpLog::kDebug + 5, "->");
}

//______________________________________________________________________________
CpFitter::~CpFitter() {
  CpDebug(CpLog::kDebug + 5, "<-");
  CpDebug(CpLog::kDebug + 5, "->");
}



void CpFitter::Print(Option_t*) {
  CpInfo(Form("Polynom Number                   : %d", fPolynomGradeNumber));
  CpInfo(Form("Number Of Active Postition Axis  : %d", GetNumberOfPositionAxis()));
  CpInfo(Form("Magentic Axis Index              : %d", fMagneticAxisIndex));
  CpInfo(Form("Axis X                           : %d [%.3f,%.3f] with %.0f bins", fActivePositionAxis[0], fMinMaxBin[0][0], fMinMaxBin[0][1], fMinMaxBin[0][2]));
  CpInfo(Form("Axis Y                           : %d [%.3f,%.3f] with %.0f bins", fActivePositionAxis[1], fMinMaxBin[1][0], fMinMaxBin[1][1], fMinMaxBin[1][2]));
  CpInfo(Form("Axis Z                           : %d [%.3f,%.3f] with %.0f bins", fActivePositionAxis[2], fMinMaxBin[2][0], fMinMaxBin[2][1], fMinMaxBin[2][2]));
  CpInfo("");
  CpInfo(Form("Output file is %s/%s", fOutputDirName.Data(), fOutputFileName.Data()));
  CpInfo("");

}


Bool_t CpFitter::Run(CpMagneticFieldMap* fieldMap) {

  // Prints current setting
  Print();

  // TODO outputfile (from config)
  fOutFile = TFile::Open(Form("%s/%s", fOutputDirName.Data(), fOutputFileName.Data()), "RECREATE");

  switch (GetNumberOfPositionAxis()) {
  case 1: Fit1D(fieldMap); break;
  case 2: Fit2D(fieldMap); break;
  case 3: Fit3D(fieldMap); break;
  default:
    CpError("Number of position axis is different then 1. Skipping ...");break;
  }

  // close file and set pointer to 0
  delete fOutFile;
  fOutFile = 0;

  // return true
  return kTRUE;
}


void CpFitter::Fit1D(CpMagneticFieldMap* fieldMap) {

  CpInfo("Fitting 1D ...");
  TStopwatch timer;
  timer.Start();
  Double_t min[3];
  Double_t max[3];
  Int_t nBins[3];
  Double_t constNum[3];
  Int_t constNumIndex[3];
  Int_t activeAxisIndex[] = {0, 0, 0};


  Int_t numMinMaxBinNum = 0;
  Int_t numConstNum = 0;
  for (Int_t i = 0;i < 3; i++) {
    if (fActivePositionAxis[i]) {
      min[numMinMaxBinNum] = fMinMaxBin[i][0];
      max[numMinMaxBinNum] = fMinMaxBin[i][1];
      nBins[numMinMaxBinNum] = (Int_t) fMinMaxBin[i][2];
      activeAxisIndex[numMinMaxBinNum] = i;
      numMinMaxBinNum++;
    } else {
      constNum[numConstNum] = fMinMaxBin[i][0];
      constNumIndex[numConstNum] = i;
      numConstNum++;
    }
  }

  TString polynomName(Form("fun_%s_1D", CpNames::AxisName(activeAxisIndex[0])));
  TF1 *myPol_1D = GetFittingFunction1D(polynomName.Data(), fPolynomGradeNumber, min[0], max[0]);
  if (!myPol_1D) {
    CpError(Form("Function %s is null. Returning ... ", polynomName.Data()));
    return;
  }

  CpDebug(CpLog::kDebug, Form("constNum[%s]=%f, constNum[%s]=%f", CpNames::AxisName(constNumIndex[0]), constNum[0], CpNames::AxisName(constNumIndex[1]), constNum[1]));

  TH1D* hist1D = new TH1D(Form("hist_%s_1D", CpNames::AxisName(activeAxisIndex[0])), Form("H_{%s} for %s=%.2f, %s=%.2f", CpNames::AxisName(fMagneticAxisIndex), CpNames::AxisName(constNumIndex[0]), constNum[0], CpNames::AxisName(constNumIndex[1]), constNum[1]), nBins[0], min[0], max[0]);
  fieldMap->FillHistogram(hist1D, fMagneticAxisIndex, activeAxisIndex[0], constNum[0], constNum[1]);
  hist1D->Fit(myPol_1D, "QN");


  if (!fOutFile) {
    CpWarning("fOutFile is null !!!");
  } else {
    CpInfo("=================================================");
    CpInfo("\t\tFitting results");
    CpInfo("=================================================");
    CpInfo("Data histogram");
    hist1D->Print();
    CpInfo("");
    CpInfo("Fit function");
    myPol_1D->Print();
    CpInfo("=================================================");
    fOutFile->cd();
    hist1D->Write();
    myPol_1D->Write();
  }

  timer.Stop();
  timer.Print();
}

void CpFitter::Fit2D(CpMagneticFieldMap* fieldMap) {

  CpInfo("Fitting 2D ...");
  TStopwatch timer;
  timer.Start();
  Double_t min[3];
  Double_t max[3];
  Int_t nBins[3];
  Double_t constNum[3];
  Int_t constNumIndex[3];
  Int_t activeAxisIndex[] = {0, 0, 0};


  Int_t numMinMaxBinNum = 0;
  Int_t numConstNum = 0;
  for (Int_t i = 0;i < 3; i++) {
    if (fActivePositionAxis[i]) {
      min[numMinMaxBinNum] = fMinMaxBin[i][0];
      max[numMinMaxBinNum] = fMinMaxBin[i][1];
      nBins[numMinMaxBinNum] = (Int_t) fMinMaxBin[i][2];
      activeAxisIndex[numMinMaxBinNum] = i;
      numMinMaxBinNum++;
    } else {
      constNum[numConstNum] = fMinMaxBin[i][0];
      constNumIndex[numConstNum] = i;
      numConstNum++;
    }
  }

  TString polynomName(Form("fun_%s%s_2D", CpNames::AxisName(activeAxisIndex[0]), CpNames::AxisName(activeAxisIndex[1])));
  TF2 *myPol_2D = GetFittingFunction2D(polynomName.Data(), fPolynomGradeNumber, min[0], max[0], min[1], max[1]);
  if (!myPol_2D) {
    CpError(Form("Function %s is null. Returning ... ", polynomName.Data()));
    return;
  }

//   CpDebug(CpLog::kDebug, Form("constNum[%s]=%f, constNum[%s]=%f", CpNames::AxisName(constNumIndex[0]), constNum[0], CpNames::AxisName(constNumIndex[1]), constNum[1]));

  TH2D* hist2D = new TH2D(Form("hist_%s%s_2D", CpNames::AxisName(activeAxisIndex[0]), CpNames::AxisName(activeAxisIndex[1])), Form("H_{%s} for %s=%.2f", CpNames::AxisName(fMagneticAxisIndex), CpNames::AxisName(constNumIndex[0]), constNum[0]), nBins[0], min[0], max[0], nBins[1], min[1], max[1]);
  fieldMap->FillHistogram(hist2D, fMagneticAxisIndex, activeAxisIndex[0], activeAxisIndex[1], constNum[0]);
  hist2D->Fit(myPol_2D, "QN");

  if (!fOutFile) {
    CpWarning("fOutFile is null !!!");
  } else {
    CpInfo("=================================================");
    CpInfo("\t\tFitting results");
    CpInfo("=================================================");
    CpInfo("Data histogram");
    hist2D->Print();
    CpInfo("");
    CpInfo("Fit function");
    myPol_2D->Print();
    CpInfo("=================================================");
    fOutFile->cd();
    hist2D->Write();
    myPol_2D->Write();
  }

  timer.Stop();
  timer.Print();
}

void CpFitter::Fit3D(CpMagneticFieldMap* fieldMap) {

  CpInfo("Fitting 3D ...");
  TStopwatch timer;
  timer.Start();
  Double_t min[3];
  Double_t max[3];
  Int_t nBins[3];
  Double_t constNum[3];
  Int_t constNumIndex[3];
  Int_t activeAxisIndex[] = {0, 0, 0};

  Int_t numMinMaxBinNum = 0;
  Int_t numConstNum = 0;
  for (Int_t i = 0;i < 3; i++) {
    if (fActivePositionAxis[i]) {
      min[numMinMaxBinNum] = fMinMaxBin[i][0];
      max[numMinMaxBinNum] = fMinMaxBin[i][1];
      nBins[numMinMaxBinNum] = (Int_t) fMinMaxBin[i][2];
      activeAxisIndex[numMinMaxBinNum] = i;
      numMinMaxBinNum++;
    } else {
      constNum[numConstNum] = fMinMaxBin[i][0];
      constNumIndex[numConstNum] = i;
      numConstNum++;
    }
  }

  TString polynomName(Form("fun_%s%s%s_3D", CpNames::AxisName(activeAxisIndex[0]), CpNames::AxisName(activeAxisIndex[1]), CpNames::AxisName(activeAxisIndex[2])));
  TF3 *myPol_3D = GetFittingFunction3D(polynomName.Data(), fPolynomGradeNumber, min[0], max[0], min[1], max[1], min[2], max[2]);
  if (!myPol_3D) {
    CpError(Form("Function %s is null. Returning ... ", polynomName.Data()));
    return;
  }

//   CpDebug(CpLog::kDebug, Form("constNum[%s]=%f, constNum[%s]=%f", CpNames::AxisName(constNumIndex[0]), constNum[0], CpNames::AxisName(constNumIndex[1]), constNum[1]));

  TH3D* hist3D = new TH3D(Form("hist_%s%s%s_3D", CpNames::AxisName(activeAxisIndex[0]), CpNames::AxisName(activeAxisIndex[1]), CpNames::AxisName(activeAxisIndex[2])), Form("H_{%s}", CpNames::AxisName(fMagneticAxisIndex)), nBins[0], min[0], max[0], nBins[1], min[1], max[1], nBins[2], min[2], max[2]);
  fieldMap->FillHistogram(hist3D, fMagneticAxisIndex);
  hist3D->Fit(myPol_3D, "QN");

  if (!fOutFile) {
    CpWarning("fOutFile is null !!!");
  } else {
    CpInfo("=================================================");
    CpInfo("\t\tFitting results");
    CpInfo("=================================================");
    CpInfo("Data histogram");
    hist3D->Print();
    CpInfo("");
    CpInfo("Fit function");
    myPol_3D->Print();
    CpInfo("=================================================");
    fOutFile->cd();
    hist3D->Write();
    myPol_3D->Write();
  }

  timer.Stop();
  timer.Print();
}

TF1* CpFitter::GetFittingFunction1D(TString name, Int_t polynomNum, Double_t min, Double_t max) {
  TF1* myPol = 0;
  Int_t nParam = CpPolynom1D::GetNumberOfParameters(polynomNum);
  CpInfo(Form("Number of parameters : %d", nParam));
  switch (polynomNum) {
  case 0:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom0, min, max, nParam);
    break;
  case 1:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom1, min, max, nParam);
    break;
  case 2:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom2, min, max, nParam);
    break;
  case 3:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom3, min, max, nParam);
    break;
  case 4:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom4, min, max, nParam);
    break;
  case 5:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom5, min, max, nParam);
    break;
  case 6:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom6, min, max, nParam);
    break;
  case 7:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom7, min, max, nParam);
    break;
  case 8:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom8, min, max, nParam);
    break;
  case 9:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom9, min, max, nParam);
    break;
  case 10:
    myPol = new TF1(name.Data(), &CpPolynom1D::DoPolynom10, min, max, nParam);
    break;
  default:
    break;
  }
  return myPol;
}

TF2* CpFitter::GetFittingFunction2D(TString name, Int_t polynomNum, Double_t min1, Double_t max1, Double_t min2, Double_t max2) {
  TF2* myPol = 0;
  Int_t nParam = CpPolynom2D::GetNumberOfParameters(polynomNum);
  CpInfo(Form("Number of parameters : %d", nParam));
  switch (polynomNum) {
  case 0:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom0, min1, max1, min2, max2, nParam);
    break;
  case 1:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom1, min1, max1, min2, max2, nParam);
    break;
  case 2:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom2, min1, max1, min2, max2, nParam);
    break;
  case 3:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom3, min1, max1, min2, max2, nParam);
    break;
  case 4:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom4, min1, max1, min2, max2, nParam);
    break;
  case 5:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom5, min1, max1, min2, max2, nParam);
    break;
  case 6:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom6, min1, max1, min2, max2, nParam);
    break;
  case 7:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom7, min1, max1, min2, max2, nParam);
    break;
  case 8:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom8, min1, max1, min2, max2, nParam);
    break;
  case 9:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom9, min1, max1, min2, max2, nParam);
    break;
  case 10:
    myPol = new TF2(name.Data(), &CpPolynom2D::DoPolynom10, min1, max1, min2, max2, nParam);
    break;
  default:
    break;
  }
  return myPol;
}

TF3* CpFitter::GetFittingFunction3D(TString name, Int_t polynomNum, Double_t min1, Double_t max1, Double_t min2, Double_t max2, Double_t min3, Double_t max3) {
  TF3* myPol = 0;
  Int_t nParam = CpPolynom3D::GetNumberOfParameters(polynomNum);
  CpInfo(Form("Number of parameters : %d", nParam));
  switch (polynomNum) {
  case 0:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom0, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 1:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom1, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 2:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom2, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 3:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom3, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 4:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom4, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 5:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom5, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 6:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom6, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 7:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom7, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 8:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom8, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 9:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom9, min1, max1, min2, max2, min3, max3, nParam);
    break;
  case 10:
    myPol = new TF3(name.Data(), &CpPolynom3D::DoPolynom10, min1, max1, min2, max2, min3, max3, nParam);
    break;
  default:
    break;
  }
  return myPol;
}

void CpFitter::SetMinMaxBin(Int_t index, Double_t min, Double_t max, Int_t bins) {
  CpDebug(CpLog::kDebug + 5, "<-");
  fMinMaxBin[index][0] = min;
  fMinMaxBin[index][1] = max;
  fMinMaxBin[index][2] = bins;
  CpDebug(CpLog::kDebug + 5, "->");
}


void CpFitter::SetActivePositionAxis(Bool_t isX, Bool_t isY, Bool_t isZ) {
  CpDebug(CpLog::kDebug + 5, "<-");
  fActivePositionAxis[0] = isX;
  fActivePositionAxis[1] = isY;
  fActivePositionAxis[2] = isZ;
  CpDebug(CpLog::kDebug + 5, "->");
}


void CpFitter::SetActivePositionAxis(Int_t index, Bool_t isActive) {
  fActivePositionAxis[index] = isActive;
}

Int_t CpFitter::GetNumberOfPositionAxis() {
  CpDebug(CpLog::kDebug + 5, "<-");
  Int_t i, num = 0;
  for (i = 0;i < 3;i++)
    if (fActivePositionAxis[i]) num++;
  CpDebug(CpLog::kDebug + 5, "->");
  return num;

}

void CpFitter::SetPositionsAxisFromString(TString str) {
  CpDebug(CpLog::kDebug + 5, "<-");
  CpDebug(CpLog::kDebug, Form("%s", str.Data()));
  TObjArray* array = str.Tokenize("][");
  TObjArray* array2 = 0;
  TString strr;
  for (Int_t i = 0;i < array->GetEntriesFast();i++) {
    strr = ((TObjString *) array->At(i))->GetString();
    if (strr.Contains(";")) {
      SetActivePositionAxis(i, kTRUE);
      array2 = strr.Tokenize(";");
      SetMinMaxBin(i, (((TObjString *) array2->At(0))->GetString()).Atof(), (((TObjString *) array2->At(1))->GetString()).Atof(), (((TObjString *) array2->At(2))->GetString()).Atof());
    } else {
      SetActivePositionAxis(i, kFALSE);
      SetMinMaxBin(i, strr.Atof(), strr.Atof(), 0);

    }

    CpDebug(CpLog::kDebug, Form("%s", strr.Data()));
  }

  CpDebug(CpLog::kDebug + 5, "->");
}
